#include "Session.h"

#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/epoll.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>

#include <iostream>
using std::cin;
using std::cout;
using std::endl;
#define likely(x) __builtin_expect(!!(x), 1)
#define unlikely(x) __builtin_expect(!!(x), 0)

static unsigned long long htonll(unsigned long long val)
{
    if(__BYTE_ORDER == __LITTLE_ENDIAN)  
    {
         return (((unsigned long long )htonl((int)((val << 32) >> 32))) << 32) | (unsigned int)htonl((int)(val >> 32));  
    }  
    else if (__BYTE_ORDER == __BIG_ENDIAN)  
    {  
         return val;  
    }  
}  


static unsigned long long ntohll(unsigned long long val)  
{  
    if (__BYTE_ORDER == __LITTLE_ENDIAN)
    {
        return (((unsigned long long )ntohl((int)((val << 32) >> 32))) << 32) | (unsigned int)ntohl((int)(val >> 32));  
    }  
    else if (__BYTE_ORDER == __BIG_ENDIAN)  
    {  
        return val;  
    }
 }

static inline void _parse_uri(const char* uri, char* ip, uint16_t* port) {
    sscanf(uri, "%[^:]:%hu", ip, port);
    cout << "ip: "<< ip <<" ,port: "<< *port << endl;
}

static inline void _get_route(char* buffer, struct RouteInf* inf)
{
    struct RouteInf* peer = (struct RouteInf*)buffer;
    inf->self_id = ntohl(peer->self_id);
    inf->r_key = ntohl(peer->r_key);
    inf->qpn = ntohl(peer->qpn);
    inf->lid = ntohs(peer->lid);
    inf->ring_begin = ntohll(peer->ring_begin);
    inf->ring_end = ntohll(peer->ring_end);
    memcpy(inf->gid, peer->gid, 16);
}

//tcp common api
static int _tcp_set_nonblocking(int sock)
{
    int flags;
    if ((flags = fcntl(sock, F_GETFL, 0)) == -1) 
    {
        cout << "Failed to set non-blocking when get socket flag" << endl;
        return -1;
    }
    if (fcntl(sock, F_SETFL, flags | O_NONBLOCK) == -1) 
    {
        cout << "Failed to set non-blocking when set socket flag" << endl;
        return -1;
    }
    return 0;
}

static int _tcp_create_socket()
{
    int sock;
    int on = 1;

    if ((sock = socket(PF_INET, SOCK_STREAM, 0)) < 0) 
    {
        cout << "Failed to create tcp socket" << endl;;
        return -1;
    }

    if ((setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on))) < 0) 
    {
        cout << "Failed to set address reuse" << endl;
        return -1;
    }

    //tcp_set_nonblocking(sock);
    return sock;
}


static int _tcp_listen(int sock, char* src_ip, uint16_t src_port)
{
    struct sockaddr_in source_addr;
    //int socket;
    //int on = 1;

    memset(&source_addr, 0, sizeof(struct sockaddr_in));
    source_addr.sin_family = AF_INET;//IPV4
    if (!src_ip) 
    {
        source_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    } else {
        inet_aton(src_ip, (struct in_addr*)&source_addr.sin_addr);
    }

    source_addr.sin_port = (src_port == -1) ? htons(7654) : htons(src_port);

    if (bind(sock, (struct sockaddr*)&source_addr, sizeof(struct sockaddr)) < 0) 
    {
        cout << "Failed to bind" << endl;
        return -1;
    }

    if (listen(sock, 6) < 0)
    {
        cout << "Failed to listen" << endl;
        return -1;
    }
    return 0;
}

/*int tcp_accept(int sock)
{
    //struct sockaddr_in remote_addr;
    //socklen_t sin_size = sizeof(struct sockaddr_in);
    //return accept(sock, (struct sockaddr*)&remote_addr, &sin_size);
    return accept(sock, NULL, NULL);
}*/

static int _tcp_connect(int sock, char* dest_ip, uint16_t dest_port)
{
    struct sockaddr_in dest_addr;
    //int socket;
    memset(&dest_addr, 0, sizeof(struct sockaddr_in));

    dest_addr.sin_family = AF_INET;//IPV4
    dest_addr.sin_port = (dest_port == -1) ? htons(7654) : htons(dest_port);
    inet_aton(dest_ip, (struct in_addr*)&dest_addr.sin_addr);

    return connect(sock, (struct sockaddr*)&dest_addr, sizeof(struct sockaddr));
}

Session::Session(MemoryManager* mm, Rdma* res) : mm_(mm), rdma_res_(res),  processed_(0), inflight_(0),credits_(32), status_(WAITING_EST)
{

}

Session::~Session()
{

}

int Session::Init(uint32_t& cur_sess_id)
{
    if(!mm_ || !rdma_res_)
    {
        cout << "Memory Manager " << mm_ << " or Rdma Resource " << rdma_res_ << " is null" << endl;
        return OTHER_ERROR;
    }
    //create qp
    qp_ = rdma_res_->RdmaCreateQueuepair();
    if(!qp_)
    {
        return CREATE_QP_FAILED;
    }
    //create ring buffer
    self_head_ = self_tail_ = self_rbuf_.begin = (uint64_t)mm_->AllocRing();
    if(!self_rbuf_.begin)
    {
        return ALLOC_RING_BUF_FAILED;
    }
    self_rbuf_.end = self_rbuf_.begin + RING_SIZE;
    
    if(cur_sess_id >= INT32_MAX)
    {
        cur_sess_id = 0;
    } 
    self_id_ = ++cur_sess_id;

    
    return OK;

}

int Session::Destroy() {
    if(mm_->FreeRing((void*)self_rbuf_.begin))
    {
        return FREE_RING_BUF_FAILED;
    }

    if(rdma_res_->RdmaDestroyQueuepair(qp_))
    {
        return DESTROY_QP_FAILED;
    }
    return OK;
}


int Session::SessConnect(const char* remote_uri) {
    char ip[20];
    uint16_t port;
    _parse_uri(remote_uri, ip, &port);

    sock_ = _tcp_create_socket();

    if(_tcp_connect(sock_, ip, port) < 0)
    {
        cout << "failed to tcp_connect\n" << endl;
        return TCP_ESTABLISHED_FAILED;
    }

    union ibv_gid gid;
    struct RouteInf self_inf;

    self_inf.self_id = htonl(self_id_);

    self_inf.r_key = htonl(rdma_res_->get_mr()->rkey);
    self_inf.qpn = htonl(qp_->qp_num);  
    self_inf.lid = htons(rdma_res_->get_lid());

    if (ibv_query_gid(rdma_res_->get_device_ctx(), rdma_res_->get_port_id(), DEFAULT_GID_INDEX, &gid)) {
        cout << "Failed to query gid" <<endl;
        close(sock_);
        return OTHER_ERROR;
    }

    memcpy(self_inf.gid, &gid, 16);

    self_inf.ring_begin = htonll(self_rbuf_.begin);
    self_inf.ring_end = htonll(self_rbuf_.end);

    int size = sizeof(struct RouteInf);
    int send_size = 0;

    if ((send_size = write(sock_, (char*)&self_inf, size)) < size) {
        cout << "Failed to send route information" << endl;;
        close(sock_);
        return WRITE_ROUTE_FAILED;
    }
    //printf("begin to recv route inf\n");
    char buffer[128];

    //struct route_inf* peer_inf = (struct route_inf*)malloc(size);
    int total_read = 0;
    while (total_read < size) {
        total_read += read(sock_, buffer + total_read, size - total_read);
    }
    struct RouteInf peer_inf;
    _get_route(buffer, &peer_inf);

    if(rdma_res_->ModifyStateINIT(qp_))
    {
        close(sock_);
        return MODIFY_QP_FAILED;
    }
    if(rdma_res_->ModifyStateRTR(qp_, peer_inf.qpn, peer_inf.lid, peer_inf.gid))
    {
        close(sock_);
        return MODIFY_QP_FAILED;
    }
    if(rdma_res_->ModifyStateRTS(qp_))
    {
        close(sock_);
        return MODIFY_QP_FAILED;
    }

    peer_id_ = peer_inf.self_id;
    peer_rkey_ = peer_inf.r_key;
    peer_rbuf_.begin = peer_inf.ring_begin;
    peer_rbuf_.end = peer_inf.ring_end;
    tail_copy_ = head_copy_ = peer_rbuf_.begin;
    is_full_ = false;
    close(sock_);
    status_ = CONNECTED;
    
    //rdma_res_->RdmaQueryQueuepair(qp_);
    return OK;
}

int Session::SessListen(const char* local_uri, int& efd) {
    char ip[20];
    uint16_t port;
    _parse_uri(local_uri, ip, &port);

    sock_ = _tcp_create_socket();
    if(sock_ < 0)
    {
        return TCP_ESTABLISHED_FAILED;
    }
    if(_tcp_listen(sock_, ip, port) < 0)
    {
        close(sock_);
        return TCP_ESTABLISHED_FAILED;
    }
    if(_tcp_set_nonblocking(sock_) < 0)
    {
        close(sock_);
        return TCP_ESTABLISHED_FAILED;
    }

    //Add to epoll, randomly set 16
    efd = epoll_create(16);

    struct epoll_event e;
    e.data.fd = sock_;
    e.events = EPOLLIN | EPOLLET;
    if(epoll_ctl(efd, EPOLL_CTL_ADD, e.data.fd, &e) < 0)
    {
        return TCP_ESTABLISHED_FAILED;
    }
    cout << "Listen Successful..." << endl;
    return OK;
}

int Session::SessAccept(int efd, std::unordered_map<uint32_t, Session*>& id2sess, MemoryManager* mm, Rdma* res, uint32_t& cur_sess_id)
{
    //char buffer[128];
    //struct RouteInf peer_inf;
    struct epoll_event events[16];
    //struct epoll_event new_event;
    struct EpollEventCtx* ctx;
    int listen_fd = sock_;

    int num = epoll_wait(efd, events, 15, 0);
    if(num < 0)
    {
        cout << "epoll wait error" << endl;
        return TCP_ESTABLISHED_FAILED;
    }
    for (int i = 0; i < num; i++) {
        if (events[i].events & (EPOLLERR | EPOLLHUP)) 
        {
            cout << "epoll error, maybe peer closed the socket" << endl;
            close(events[i].data.fd);
            continue;
        }
        else if (events[i].events & EPOLLIN) 
        {
            cout << "get input event" <<endl;
            if (events[i].data.fd == listen_fd) 
            {
                int new_fd = -1;
                 
                while ((new_fd = accept(listen_fd, NULL, NULL)) != -1) 
                {
                    new_fd = accept(listen_fd, NULL, NULL);
                    //cout << "failed to accept, try again" << endl;
                
                    if (_tcp_set_nonblocking(new_fd) < 0) 
                    {
                        cout << "socket " << new_fd << " failed to set nonblocking..." << endl;
                        return TCP_ESTABLISHED_FAILED;
                    }
                    struct epoll_event new_event;
                    new_event.data.fd = new_fd;
                    new_event.events = EPOLLIN | EPOLLET;

                    if(epoll_ctl(efd, EPOLL_CTL_ADD, new_fd, &new_event) < 0)
                    {
                        cout << "epoll ctl failed" << endl;
                        return TCP_ESTABLISHED_FAILED;
                    }
                    cout << "accept a new connection, fd = " << new_fd <<endl;
                }
            }
            else {
                Session* new_sess = new Session(mm, res);
                if(new_sess->Init(cur_sess_id) < 0)
                {
                    delete new_sess;
                    return TCP_ESTABLISHED_FAILED;
                }
                

                if(new_sess->SessReadRoute(events[i].data.fd) < 0)
                {
                    delete new_sess;
                    return TCP_ESTABLISHED_FAILED;
                }

                ctx = new struct EpollEventCtx;


                ctx->fd = events[i].data.fd;
                ctx->sess = new_sess;
                struct epoll_event new_event;

                new_event.data.ptr = (void*)ctx;

                new_event.events = EPOLLOUT | EPOLLET;

                if(epoll_ctl(efd, EPOLL_CTL_MOD, events[i].data.fd, &new_event)<0)
                {
                    delete new_sess;
                    delete ctx;
                    return TCP_ESTABLISHED_FAILED;
                }

            }
        }
        else if ((events[i].events & EPOLLOUT) && events[i].data.fd != listen_fd) {
            ctx = (struct EpollEventCtx*)events[i].data.ptr;
            if (ctx->fd != listen_fd) 
            {
                //printf("begin to send route inf, fd = %d.....\n", events[i].data.fd);
                Session* new_sess = ctx->sess;
                if(new_sess->SessWriteRoute(ctx->fd) < 0)
                {
                    delete new_sess;
                    delete ctx;
                    return TCP_ESTABLISHED_FAILED;
                }
                cout << "Client " << new_sess->self_id_ <<" connect successful..." << endl;
                id2sess[new_sess->self_id_] = new_sess;
                close(ctx->fd);
                new_sess->status_ = CONNECTED;
            }
            else{
                return TCP_ESTABLISHED_FAILED;

            }
        }
        else
        {
            cout << "Error epll events" << endl;
            return TCP_ESTABLISHED_FAILED;

        }
    }
    return OK;

}

int Session::SessReadRoute(int fd) 
{
    uint32_t total_size = 0;
    uint32_t read_size = 0;
    struct RouteInf peer_inf;
    char buffer[64];
    sock_ = fd;
    while (total_size < sizeof(struct RouteInf)) {
        read_size = read(sock_, buffer + total_size, sizeof(struct RouteInf) - total_size);
        if(read_size < 0)
        {
            cout << "socket write errno = " << errno << endl;
            return READ_ROUTE_FAILED;
        } else {
            total_size += read_size;
        }
    }
    _get_route(buffer, &peer_inf);
    //printf("peer_qp_num : %d, %d, peer_id : %d\n", peer_inf.qpn, peer_inf.qpn2, peer_inf.sess_id);

    if(rdma_res_->ModifyStateINIT(qp_))
    {
        return MODIFY_QP_FAILED;
    }
    if(rdma_res_->ModifyStateRTR(qp_, peer_inf.qpn, peer_inf.lid, peer_inf.gid))
    {
        return MODIFY_QP_FAILED;
    }
    if(rdma_res_->ModifyStateRTS(qp_))
    {
        return MODIFY_QP_FAILED;
    }
    peer_rkey_ = peer_inf.r_key;
    peer_id_ = peer_inf.self_id;
    peer_rbuf_.begin = peer_inf.ring_begin;
    peer_rbuf_.end = peer_inf.ring_end;
    tail_copy_ = head_copy_ = peer_rbuf_.begin;
    is_full_ = false;
    
    return OK;
}

int Session::SessWriteRoute(int fd) {
    //int rc = 0;
    union ibv_gid gid;
    struct RouteInf self_inf;
    if(fd != sock_)
    {
        cout << "read route use socket " << sock_ << ", but write use socket " << fd << endl;
        return OTHER_ERROR;
    }
    //set route information locally
    self_inf.self_id = htonl(self_id_);
    self_inf.r_key = htonl(rdma_res_->get_mr()->rkey);
    self_inf.qpn = htonl(qp_->qp_num);  
    self_inf.lid = htons(rdma_res_->get_lid());

    if (ibv_query_gid(rdma_res_->get_device_ctx(), rdma_res_->get_port_id(), DEFAULT_GID_INDEX, &gid)) {
        cout << "Failed to query gid" <<endl;
        return OTHER_ERROR;
    }

    memcpy(self_inf.gid, &gid, 16);

    self_inf.ring_begin = htonll(self_rbuf_.begin);
    self_inf.ring_end = htonll(self_rbuf_.end);

    int size = sizeof(struct RouteInf);
    int total_size = 0;
    int write_size = 0;
    while (total_size < size) {
        write_size = write(fd, (char*)&self_inf + total_size, size - total_size);
        if (write_size < 0) {
            cout << "socket write errno = " << errno << endl;
            return WRITE_ROUTE_FAILED;
        } else {
            total_size += write_size;
        }
    }
    return OK;
}

uint64_t Session::get_peer_buf(uint64_t len) 
{
    uint64_t alloc = 0;
    if(is_full_ || len == 0){
        //std::cout << "tail = " << tail_copy_ << ", head = " << head_copy_ <<std::endl;
        return alloc;
    }
    if(tail_copy_ >= head_copy_)
    {
        alloc = tail_copy_;
        tail_copy_ = (tail_copy_ < peer_rbuf_.end - len) ? tail_copy_ + len : peer_rbuf_.begin;
    }
    else
    {
        if(tail_copy_ <= head_copy_ - len)
        {
            alloc = tail_copy_;
            tail_copy_ += len;
        }
    }
    is_full_ = (tail_copy_ == head_copy_);

    return alloc;
}

void Session::free_peer_buf(uint32_t len)
{
    //buf is empty
    if(!is_full_ && tail_copy_ == head_copy_) return;
    if(head_copy_ < peer_rbuf_.end - len)
    {
        head_copy_+=len;
    }
    else if(head_copy_ >= peer_rbuf_.end - len)
    {
        head_copy_ = peer_rbuf_.begin;
    }
    if(is_full_ && len) is_full_ = false;
}

uint64_t Session::poll_msg()
{
    
    volatile uint32_t *p = (volatile uint32_t*)self_tail_;
    if(*p){
        //uint32_t msg_len = (uint32_t)*p;
        //garbage data
        if(self_tail_ + (*p) >= self_rbuf_.end + MSG_SIZE) return 0;
        volatile uint32_t *canary = (volatile uint32_t*)((char*)p + (*p));
        if((*canary) && (*(p+1)) && (*canary) == ((*p)+(*(p+1)))) 
        {
            //cout  << "[recv msg] msg_len : " << *p << " sync : " << (*(p+1)) << " canary : " << *canary << endl; 
            //prevent poll it again before call rpc_handle_complete
            *canary = 0;
            if(self_tail_ < self_rbuf_.end - (*p) - 4)
            {

                self_tail_ = self_tail_ + (*p) + 4; 
                
            }
            else
            {
                self_tail_ = self_rbuf_.begin;
            }
            
            //std::cout << "!!!POLL!!! [" << self_rbuf_.begin << " , " << self_rbuf_.end << "]" << " recv" << (*p)+4 << "\nhead = "
            //<< self_head_ << "\ntail = " << self_tail_ << std::endl;
            return (uint64_t)p;
        }
        else
        { 
            return 0;
        }
    }
    else 
    {
        return 0;
    }
}


uint32_t Session::free_msg(uint64_t msg)
{
    uint32_t total_free = 0;
    volatile uint32_t *p = (uint32_t*)msg;
    
    uint32_t msg_len = (uint32_t)*p;

    //switches to canary
    volatile uint64_t *q = (uint64_t*)((char*)p + 8);
    
    //set signature 1
    *q = 0x8888777799995555ull;

    p = (volatile uint32_t*)self_head_;

    //point to second uint32_t of msg, if freed, must be zero
    q = (volatile uint64_t*)((char*)p + 8);
    while((self_head_ + (*p) < self_rbuf_.end + MSG_SIZE)&&(*p) && (*q) == 0x8888777799995555ull)
    {
        //must clear canary
        *q = 0ull; 
        //std::cout << "expected : " << (uint32_t)(prev_free_sync_+1) << "give : " << *sync<< std::endl;

        uint32_t msg_add_canary = (*p) + 4;
        total_free += msg_add_canary;
        
        //no more than end of ring buf, move self_head to end of cannary
        if(self_head_ < self_rbuf_.end - msg_add_canary)
        {
            self_head_ = self_head_ + msg_add_canary;
            
        }
        //otherwise, move to begin of ring buf
        else
        {
            self_head_ = self_rbuf_.begin;
            break;
        }
         

        p = (volatile uint32_t*)self_head_;

        q = (volatile uint64_t*)((char*)p + 8);
    }
    //if(total_free) memset(clear_begin, 0, total_free);

    //std::cout << "!!!FREE!!! [" << self_rbuf_.begin << " , " << self_rbuf_.end << "]" << " free" << total_free << "\nhead = "
    //<< self_head_ << "\ntail = " << self_tail_ << std::endl;
    //if(self_head_ == self_tail_) 
    //{
        /*std::cout << "!!!begin = " << *(uint32_t*)self_rbuf_.begin << std::endl;

        std::cout << "!!!!head == tail!!!!" << std::endl;
        uint32_t *_p = (uint32_t*)self_head_;
    
        uint32_t _msg_len = (uint32_t)*_p;

        //switches to canary
        uint64_t *_q = (uint64_t*)((char*)_p + 8);
        std::cout << "[head] msg len = " << _msg_len << " ,sign = "<< *_q <<std::endl;
*/
        

    //}

    
    return total_free;
}




